package com.zh.agent;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.security.ProtectionDomain;
import java.util.LinkedList;
import java.util.List;

import com.sun.tools.attach.VirtualMachine;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.LoaderClassPath;

public class Agent implements ClassFileTransformer {

    private static final String targetClassName = "com.zh.agent.app.Hello";
    private static final String targetMethodName = "sayHello";

    private String targetVMClassName;

    public Agent(String className) {
        this.targetVMClassName = new String(className).replaceAll("\\.", "\\/");
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

        if (!className.equals(this.targetVMClassName)) {
            // System.out.println("Not Match not transform");
            return classfileBuffer;
        }

        System.out.println("Match transform  class = " + className);

        ClassPool classPool = ClassPool.getDefault();
        CtClass cls;
        try {
            cls = classPool.get(targetClassName);
            CtMethod[] ctMethod = cls.getDeclaredMethods(targetMethodName);
            for (CtMethod method : ctMethod) {
                method.insertBefore("System.out.println(\"Agent insert before\");");
                method.insertAfter("System.out.println(\"Agent insert after\");");
                return cls.toBytecode();
            }
        } catch (Throwable e) {
            e.printStackTrace();
        }
        return classfileBuffer;
    }

    public static void premain(String args, Instrumentation inst) {
        System.out.println("Premain Start");
        inst.addTransformer(new Agent(targetClassName));
        System.out.println(args);
    }

    public static void agentmain(String args, Instrumentation inst) {
        System.out.println("Load Agent After Main run.args = " + args);

        inst.addTransformer(new ClassFileTransformer() {

            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                    ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

                System.out.println("Begin transform");
                ClassPool classPool = ClassPool.getDefault();
                classPool.appendClassPath(new LoaderClassPath(loader));
                try {
                    String jvmclassName = className.replace("/", ".");
                    CtClass clz = classPool.get(jvmclassName);
                    CtMethod[] methods = clz.getDeclaredMethods();
                    System.out.println("Load Methods count == "+methods.length);
                    for (CtMethod method : methods) {
                        System.out.println(method.getName());
                        if (method.getName().startsWith("sayHello")) {
                            method.insertBefore("System.out.println(\"before " + method.getName() + " execute!!!\");");
                            method.insertAfter("System.out.println(\"after " + method.getName() + " execute!!!\");");
                        }
                    }
                    byte[] r = clz.toBytecode();
                    clz.detach();
                    return r;
                } catch (Throwable e) {
                    e.printStackTrace();
                }

                return new byte[0];
            }
        }, true);
        for (Class cls : inst.getAllLoadedClasses()) {
            if ("com.zh.agent.app.Hello".equals(cls.getName())) {
                try {
                    System.out.println("retransform class == " + cls.getName());
                    inst.retransformClasses(cls);
                } catch (UnmodifiableClassException e) {
                    e.printStackTrace();
                }
            }
        }
        System.out.println("Attach Class Running finish");
    }

    public static void main(String[] args) {
        try {
            if (args.length < 1) {
                System.out.println("Need Pid To Attach!");
                System.exit(0);
            }
            String pid = args[0];
            VirtualMachine vm = VirtualMachine.attach(pid);
            vm.loadAgent("./build/libs/agent.jar");
            vm.detach();
        } catch (Throwable e) {
            System.out.println("Exception : " + e.getMessage());
            e.printStackTrace();
        }
    }
}
